package com.learn.security.oauth2.auth.server.config.oauth2;

import cn.hutool.core.collection.CollUtil;
import org.springframework.context.support.MessageSourceAccessor;
import org.springframework.security.authentication.AccountExpiredException;
import org.springframework.security.authentication.BadCredentialsException;
import org.springframework.security.authentication.CredentialsExpiredException;
import org.springframework.security.authentication.DisabledException;
import org.springframework.security.authentication.LockedException;
import org.springframework.security.authentication.UsernamePasswordAuthenticationToken;
import org.springframework.security.authentication.dao.DaoAuthenticationProvider;
import org.springframework.security.core.Authentication;
import org.springframework.security.core.AuthenticationException;
import org.springframework.security.core.SpringSecurityMessageSource;
import org.springframework.security.core.userdetails.UserDetailsService;
import org.springframework.security.core.userdetails.UsernameNotFoundException;
import org.springframework.security.crypto.keygen.Base64StringKeyGenerator;
import org.springframework.security.crypto.keygen.StringKeyGenerator;
import org.springframework.security.oauth2.core.AuthorizationGrantType;
import org.springframework.security.oauth2.core.ClientAuthenticationMethod;
import org.springframework.security.oauth2.core.OAuth2AccessToken;
import org.springframework.security.oauth2.core.OAuth2AuthenticationException;
import org.springframework.security.oauth2.core.OAuth2Error;
import org.springframework.security.oauth2.core.OAuth2ErrorCodes;
import org.springframework.security.oauth2.core.OAuth2RefreshToken;
import org.springframework.security.oauth2.core.endpoint.OAuth2ParameterNames;
import org.springframework.security.oauth2.jose.jws.SignatureAlgorithm;
import org.springframework.security.oauth2.jwt.JwsHeader;
import org.springframework.security.oauth2.jwt.Jwt;
import org.springframework.security.oauth2.jwt.JwtClaimNames;
import org.springframework.security.oauth2.jwt.JwtClaimsSet;
import org.springframework.security.oauth2.jwt.JwtEncoder;
import org.springframework.security.oauth2.jwt.JwtEncoderParameters;
import org.springframework.security.oauth2.server.authorization.OAuth2Authorization;
import org.springframework.security.oauth2.server.authorization.OAuth2AuthorizationService;
import org.springframework.security.oauth2.server.authorization.OAuth2TokenType;
import org.springframework.security.oauth2.server.authorization.authentication.OAuth2AccessTokenAuthenticationToken;
import org.springframework.security.oauth2.server.authorization.authentication.OAuth2ClientAuthenticationToken;
import org.springframework.security.oauth2.server.authorization.client.RegisteredClient;
import org.springframework.security.oauth2.server.authorization.context.AuthorizationServerContext;
import org.springframework.security.oauth2.server.authorization.context.AuthorizationServerContextHolder;
import org.springframework.security.oauth2.server.authorization.token.JwtEncodingContext;
import org.springframework.security.oauth2.server.authorization.token.OAuth2TokenCustomizer;
import org.springframework.util.CollectionUtils;

import java.security.Principal;
import java.time.Duration;
import java.time.Instant;
import java.util.Base64;
import java.util.Collections;
import java.util.HashMap;
import java.util.Locale;
import java.util.Map;
import java.util.Objects;
import java.util.Set;

/**
 * 处理用户名密码授权
 *
 * @author knight
 */
public class OAuth2PasswordAuthenticationProvider extends DaoAuthenticationProvider {

	private final JwtEncoder jwtEncoder;

	private final OAuth2AuthorizationService authorizationService;

	private final StringKeyGenerator refreshTokenGenerator = new Base64StringKeyGenerator(
			Base64.getUrlEncoder().withoutPadding(), 96);

	private final OAuth2TokenCustomizer<JwtEncodingContext> jwtCustomizer = (context) -> {
	};

	private final MessageSourceAccessor messages = new MessageSourceAccessor(new SpringSecurityMessageSource(),
			Locale.CHINA);

	public OAuth2PasswordAuthenticationProvider(UserDetailsService userDetailsService, JwtEncoder jwtEncoder,
			OAuth2AuthorizationService authorizationService) {
		super.setUserDetailsService(userDetailsService);
		this.jwtEncoder = jwtEncoder;
		this.authorizationService = authorizationService;
	}

	@Override
	public Authentication authenticate(Authentication authentication) throws AuthenticationException {
		// 用户名密码token
		OAuth2PasswordAuthenticationToken passwordAuthenticationToken = (OAuth2PasswordAuthenticationToken) authentication;

		// 获取客户端
		OAuth2ClientAuthenticationToken clientPrincipal = null;
		if (OAuth2ClientAuthenticationToken.class.isAssignableFrom(authentication.getPrincipal().getClass())) {
			clientPrincipal = (OAuth2ClientAuthenticationToken) authentication.getPrincipal();
		}
		if (clientPrincipal == null || !clientPrincipal.isAuthenticated()
				|| clientPrincipal.getRegisteredClient() == null) {
			throw new OAuth2AuthenticationException(OAuth2ErrorCodes.INVALID_CLIENT);
		}
		RegisteredClient registeredClient = clientPrincipal.getRegisteredClient();

		// 校验该客户端是否有密码授权类型
		if (!CollUtil.contains(registeredClient.getAuthorizationGrantTypes(), AuthorizationGrantType.PASSWORD)) {
			throw new OAuth2AuthenticationException(OAuth2ErrorCodes.UNAUTHORIZED_CLIENT);
		}

		UsernamePasswordAuthenticationToken usernamePasswordAuthenticationToken = new UsernamePasswordAuthenticationToken(
				passwordAuthenticationToken.getUsername(), passwordAuthenticationToken.getPassword());
		// 调用AbstractUserDetailsAuthenticationProvider.authenticate()进行校验
		Authentication usernamePasswordAuthentication;
		try {
			usernamePasswordAuthentication = super.authenticate(usernamePasswordAuthenticationToken);
		}
		catch (AuthenticationException e) {
			throw convertLoginExceptionToOauth2Exception(authentication, e);
		}

		Set<String> authorizedScopes = registeredClient.getScopes();

		// 生成访问令牌jwt
		Jwt jwtAccessToken = accessTokenJwt(registeredClient, usernamePasswordAuthentication,
				usernamePasswordAuthenticationToken);

		// 生成访问令牌
		OAuth2AccessToken accessToken = new OAuth2AccessToken(OAuth2AccessToken.TokenType.BEARER,
				jwtAccessToken.getTokenValue(), jwtAccessToken.getIssuedAt(), jwtAccessToken.getExpiresAt(),
				authorizedScopes);

		Map<String, Object> claims = jwtAccessToken.getClaims();

		OAuth2RefreshToken refreshToken = null;
		if (registeredClient.getAuthorizationGrantTypes().contains(AuthorizationGrantType.REFRESH_TOKEN) &&
		// 不向公共客户端发送刷新令牌
				!clientPrincipal.getClientAuthenticationMethod().equals(ClientAuthenticationMethod.NONE)) {
			refreshToken = generateRefreshToken(registeredClient.getTokenSettings().getRefreshTokenTimeToLive());
		}

		// @formatter:off
		OAuth2Authorization.Builder authorizationBuilder = OAuth2Authorization.withRegisteredClient(registeredClient)
				.principalName(usernamePasswordAuthentication.getName())
				.authorizationGrantType(AuthorizationGrantType.PASSWORD)
				.token(accessToken,
						(metadata) ->
								metadata.put(OAuth2Authorization.Token.CLAIMS_METADATA_NAME, claims))
				.authorizedScopes(authorizedScopes)
				.attribute(Principal.class.getName(), usernamePasswordAuthentication);
		// @formatter:on
		if (refreshToken != null) {
			authorizationBuilder.refreshToken(refreshToken);
		}

		OAuth2Authorization authorization = authorizationBuilder.build();
		this.authorizationService.save(authorization);

		Map<String, Object> tokenAdditionalParameters = new HashMap<>();
		claims.forEach((key, value) -> {
			if (!key.equals(OAuth2ParameterNames.SCOPE) && !key.equals(JwtClaimNames.IAT)
					&& !key.equals(JwtClaimNames.EXP) && !key.equals(JwtClaimNames.NBF)) {
				tokenAdditionalParameters.put(key, value);
			}
		});

		return new OAuth2AccessTokenAuthenticationToken(registeredClient, clientPrincipal, accessToken, refreshToken,
				tokenAdditionalParameters);

	}

	@Override
	public boolean supports(Class<?> authentication) {
		return OAuth2PasswordAuthenticationToken.class.isAssignableFrom(authentication);
	}

	/**
	 * 登录异常转换为oauth2异常
	 * @param authentication 身份验证
	 * @param authenticationException 身份验证异常
	 * @return {@link OAuth2AuthenticationException}
	 */
	private OAuth2AuthenticationException convertLoginExceptionToOauth2Exception(Authentication authentication,
			AuthenticationException authenticationException) {
		if (authenticationException instanceof UsernameNotFoundException) {
			return new OAuth2AuthenticationException(new OAuth2Error(OAuth2ErrorCodesExpand.USERNAME_NOT_FOUND,
					this.messages.getMessage("JdbcDaoImpl.notFound", new Object[] { authentication.getName() },
							"Username {0} not found"),
					""));
		}
		if (authenticationException instanceof BadCredentialsException) {
			return new OAuth2AuthenticationException(
					new OAuth2Error(OAuth2ErrorCodesExpand.BAD_CREDENTIALS, this.messages.getMessage(
							"AbstractUserDetailsAuthenticationProvider.badCredentials", "Bad credentials"), ""));
		}
		if (authenticationException instanceof LockedException) {
			return new OAuth2AuthenticationException(new OAuth2Error(OAuth2ErrorCodesExpand.USER_LOCKED, this.messages
					.getMessage("AbstractUserDetailsAuthenticationProvider.locked", "User account is locked"), ""));
		}
		if (authenticationException instanceof DisabledException) {
			return new OAuth2AuthenticationException(new OAuth2Error(OAuth2ErrorCodesExpand.USER_DISABLE,
					this.messages.getMessage("AbstractUserDetailsAuthenticationProvider.disabled", "User is disabled"),
					""));
		}
		if (authenticationException instanceof AccountExpiredException) {
			return new OAuth2AuthenticationException(new OAuth2Error(OAuth2ErrorCodesExpand.USER_EXPIRED, this.messages
					.getMessage("AbstractUserDetailsAuthenticationProvider.expired", "User account has expired"), ""));
		}
		if (authenticationException instanceof CredentialsExpiredException) {
			return new OAuth2AuthenticationException(new OAuth2Error(OAuth2ErrorCodesExpand.CREDENTIALS_EXPIRED,
					this.messages.getMessage("AbstractUserDetailsAuthenticationProvider.credentialsExpired",
							"User credentials have expired"),
					""));
		}
		return new OAuth2AuthenticationException(OAuth2ErrorCodesExpand.UN_KNOW_LOGIN_ERROR);
	}

	/**
	 * 生成刷新令牌
	 * @param tokenTimeToLive 令牌时间生活
	 * @see org.springframework.security.oauth2.server.authorization.token.OAuth2RefreshTokenGenerator
	 * @return {@link OAuth2RefreshToken}
	 */
	private OAuth2RefreshToken generateRefreshToken(Duration tokenTimeToLive) {
		Instant issuedAt = Instant.now();
		Instant expiresAt = issuedAt.plus(tokenTimeToLive);
		return new OAuth2RefreshToken(this.refreshTokenGenerator.generateKey(), issuedAt, expiresAt);
	}

	/**
	 * 访问令牌jwt
	 * @see org.springframework.security.oauth2.server.authorization.token.OAuth2AccessTokenGenerator
	 * @param registeredClient 注册客户端
	 * @param usernamePasswordAuthentication 用户名密码认证
	 * @param usernamePasswordAuthenticationToken 用户名密码身份验证令牌
	 * @return {@link Jwt}
	 */
	private Jwt accessTokenJwt(RegisteredClient registeredClient, Authentication usernamePasswordAuthentication,
			UsernamePasswordAuthenticationToken usernamePasswordAuthenticationToken) {

		Instant issuedAt = Instant.now();
		Instant expiresAt = issuedAt.plus(registeredClient.getTokenSettings().getAccessTokenTimeToLive());
		// @formatter:off
		JwsHeader.Builder headersBuilder = JwsHeader.with(SignatureAlgorithm.RS256);
		JwtClaimsSet.Builder claimsBuilder = JwtClaimsSet.builder();

		AuthorizationServerContext authorizationServerContext = AuthorizationServerContextHolder.getContext();
		if (Objects.nonNull(authorizationServerContext)) {
			// 认证信息者的唯一标识(授权服务的地址)
			claimsBuilder.issuer(authorizationServerContext.getIssuer());
		}
		claimsBuilder
				// 用户
				.subject(usernamePasswordAuthenticationToken.getName())
				// clientId
				.audience(Collections.singletonList(registeredClient.getClientId()))
				// 发布时间
				.issuedAt(issuedAt)
				// 过期时间
				.expiresAt(expiresAt)
				.notBefore(issuedAt);
		if (!CollectionUtils.isEmpty(registeredClient.getScopes())) {
			// scope
			claimsBuilder.claim(OAuth2ParameterNames.SCOPE, registeredClient.getScopes());
		}
		// @formatter:on

		if (this.jwtCustomizer != null) {
			// @formatter:off
			JwtEncodingContext context = JwtEncodingContext.with(headersBuilder, claimsBuilder)
					.registeredClient(registeredClient)
					.principal(usernamePasswordAuthentication)
					.authorizedScopes(registeredClient.getScopes())
					.tokenType(OAuth2TokenType.ACCESS_TOKEN)
					.authorizationGrantType(AuthorizationGrantType.PASSWORD)
					.authorizationGrant(usernamePasswordAuthenticationToken)
					.build();
			// @formatter:on
			this.jwtCustomizer.customize(context);
		}

		return this.jwtEncoder.encode(JwtEncoderParameters.from(headersBuilder.build(), claimsBuilder.build()));
	}

}
